
local tokenRate = tonumber(ARGV[1])
local bucketSize = tonumber(ARGV[2])
local currentTime = tonumber(ARGV[3])
local tokensNeeded = tonumber(ARGV[4])

local timeToFillBucket = bucketSize/tokenRate
local expiryTime = math.floor(timeToFillBucket*2)

local remainingTokens = tonumber(redis.call("get", KEYS[1]))
if remainingTokens == nil then
    remainingTokens = bucketSize
end

local lastTokenUpdateTime = tonumber(redis.call("get", KEYS[2]))
if lastTokenUpdateTime == nil then
    lastTokenUpdateTime = 0
end

local passedTime = math.max(0, currentTime-lastTokenUpdateTime)
local totalTokens = math.min(bucketSize, remainingTokens+(passedTime*tokenRate))

local requestAllowed = totalTokens >= tokensNeeded
local updatedTokens = totalTokens
if requestAllowed then
    updatedTokens = totalTokens - tokensNeeded
end

redis.call("setex", KEYS[1], expiryTime, updatedTokens)
redis.call("setex", KEYS[2], expiryTime, currentTime)

return requestAllowed